# -*- coding:utf-8 -*-
import numpy as np
import time


class Ising():
    """
    A class to implement Ising models.

    ...

    Attributes
    ----------
    J : numpy.ndarray
        spin interactions matrix (must be semi-definite positive)
    h : numpy.ndarray
        magnectic field effect vector
    dimension : int
        number of spins
    ground_state : numpy.ndarray
        vector of spins orientation to minimize the energy
    optimization_logs : dict
        data about the optimization of the model
    """

    def __init__(self, J: np.ndarray, h: np.ndarray, assert_parameters: bool = True) -> None:

        """
        Constructs all the necessary attributes for the Ising object.

        Parameters
        ----------
            J : numpy.ndarray
                spin interactions matrix (must be semi-definite positive)
            h : numpy.ndarray
                magnectic field effect vector
            assert_parameters : bool, optional
                check the format of the inputs (default is True)
        """

        self.J = J
        self.h = h
        self.e = 0

        self.dimension = J.shape[0]

        self.ground_state = np.ones((self.dimension, 1))
        self.optimization_logs = dict()

        # Check parameters

        if assert_parameters:
            self.__assert__()

    def __assert__(self) -> None:

        """
        Checks the format of the attributes.

        Returns
        -------
        float
        """

        # Checking types
        assert isinstance(self.J,
                          np.ndarray), f"WARNING: The type of J must be a numpy array, instead got {type(self.J)}"
        assert isinstance(self.h,
                          np.ndarray), f"WARNING: The type of h must be a numpy array, instead got {type(self.h)}"

        # Checking dimensions
        assert self.J.shape[0] == self.J.shape[1], f"WARNING: J must be a square matrix, instead got {self.J.shape}"
        assert self.h.shape[0] == self.J.shape[
            0], f"WARNING: The dimension of h must fits J's, instead of {self.J.shape[0]} got {self.h.shape[0]}"
        assert self.h.shape[
                   1] == 1, f"WARNING: h must be a column vector with dimensions of this pattern: (n,1), instead got {self.h.shape}"

        # Checking J's properties
        assert np.allclose(self.J, self.J.T), "WARNING: J must be symmetric"
        # assert not np.any(self.J == np.zeros(self.J.shape)), "WARNING: J must not have null elements"

    def energy(self) -> float:

        """
        Computes the Ising energy of the model.

        Returns
        -------
        float
        """

        if self.ground_state is None:
            return None

        else:
            energy = -0.5 * self.ground_state.T @ self.J @ self.ground_state + self.ground_state.T @ self.h
            self.e = energy[0][0].copy()
            return energy[0][0]

    def get_ground_state(
            self,
            spin
    ) -> None:
        if (np.sign(spin).any()):
            self.ground_state = np.sign(spin)


class SBModel():
    """
    A class to implement Ising problems.
    """

    def __to_Ising__(self) -> Ising:
        """
        Generate the equivalent Ising model of the problem.

        Returns
        -------
        Ising
        """

        pass

    def __from_Ising__(self, ising: Ising) -> None:
        """
        Retrieves information from the optimized equivalent Ising model.
        Modifies the object's attributes in place.

        Parameters
        ----------
            ising : Ising
                equivalent Ising model of the problem

        Returns
        -------
        None
        """

        pass


class MaxCut(SBModel):
    """
    A representation of the Markowitz model for portolio optimization.
    Portfolio only takes integer stocks.

    ...

    Attributes
    ----------
    covariance : numpy.ndarray
        the correlation matrix between the assets
    expected_return : numpy.ndarray
        expected return on the investment
    risk_coefficient : float
        risk aversion on the investment
    assets_list : list
        list of the assets to invest in
    number_of_assets : int
        number of different assets
    number_of_bits : int
        number of bits for the binary decomposition of the assets stocks
    portfolio : numpy.ndarray
        array of stocks to purchase per asset
    M : numpy.ndarray
        integer-binary conversion matrix
    ones : numpy.ndarray
        ones vector
    """

    @classmethod
    def from_gset(
            cls,
            gset_path: str
    ):

        """
        Retrieves the data for the Markowitz model from .csv files.
        Only works with the files in the ./data folder.
        """

        with open(gset_path, mode='r') as f:
            title = f.readline()
            if title == 'ROWS':
                print("It is eCut problem.")
                # TODO: process the eCut problem.
                exit(0)
            else:
                vertex_str, edge_str = title.split()
                vertex = int(vertex_str)
                edge = int(edge_str)
                weight = np.zeros((vertex, vertex))
                edge_relationship = f.readlines()
                for cnt in range(edge):
                    edge_a_str, edge_b_str, weight_a_b_str = edge_relationship[cnt].split()
                    edge_a = int(edge_a_str)
                    edge_b = int(edge_b_str)
                    weight_a_b = int(weight_a_b_str)
                    weight[edge_a - 1][edge_b - 1] = weight_a_b
                    weight[edge_b - 1][edge_a - 1] = weight_a_b

        return MaxCut(
            vertex,
            edge,
            weight
        )

    def __init__(
            self,
            vertex: int,
            edge: int,
            weight: np.ndarray,
            assert_parameters: int = 1
    ) -> None:

        """
        Constructs all the necessary attributes for the Markowitz object.

        Parameters
        ----------
            vertex : int
                the vertex num of
            expected_return : numpy.ndarray
                expected return on the investment
            risk_coefficient : float
                risk aversion on the investment
            assets_list : list
                list of the assets to invest in
            number_of_bits : int
                number of bits for the binary decomposition of the assets stocks
            assert_parameters : bool, optional
                check the format of the inputs (default is True)
        """

        # Data
        self.vertex = vertex
        self.edge = edge
        self.weight = weight
        self.cut_value = vertex
        if assert_parameters:
            self.__assert__()

    def __repr__(self) -> str:

        return f"""Utility gain: {self.utlity_function()}
        {self.as_dataframe()}
        """

    def __assert__(self) -> None:

        """
        Checks the format of the attributes.

        Returns
        -------
        float
        """

        # Checking types

        # Checking dimensions

        pass

    def __to_Ising__(self) -> Ising:

        J = -self.weight
        h = np.zeros((self.vertex, 1))

        return Ising(J, h)

    def __from_Ising__(self, ising: Ising) -> None:
        if self.cut_value <= -0.25 * np.sum(
                -self.weight) - 0.25 * ising.ground_state.T @ self.weight @ ising.ground_state:
            self.lowest_ground_state = ising.ground_state.copy()
            self.cut_value = -0.25 * np.sum(
                -self.weight) - 0.25 * ising.ground_state.T @ self.weight @ ising.ground_state


class sbpso():
    """
    The instance of the SBPSO Algorithm, the PSO Algorithm combined with Simulated Bifurcation Algorithm.

    """

    def __init__(
            self,
            ising: Ising,
            weight: float,
            lr: tuple,
            maxgen: int,
            sizepop: int,
            rangepop: tuple,
            rangespeed: tuple,
            detuning_frequency: float = 1,
            kerr_constant: float = 1,
            pressure=lambda t: 0.01 * t,
            time_step: float = 0.25,
            symplectic_parameter: int = 1,
            convergence_threshold: int = 35,
            sampling_period: int = 60,
    ) -> None:
        self.ising = ising  # 伊辛模型
        self.weight = weight  # 惯性权重
        self.rangespeed = rangespeed  # 粒子的速度范围限制
        self.rangepop = rangepop  # 粒子的位置的范围限制
        self.sizepop = sizepop  # 种群规模
        self.maxgen = maxgen  # 最大迭代次数
        self.lr = lr  # 加速常数

        self.dimpop = ising.dimension  # 粒子维度
        self.x = np.zeros((self.dimpop, self.sizepop))  # 粒子位置
        self.v = np.zeros((self.dimpop, self.sizepop))  # 粒子速度
        self.pbestpop = np.zeros((self.dimpop, self.sizepop))  # 各粒子历史最优位置
        self.pbestval = np.zeros((1, self.sizepop)) + np.inf  # 各粒子历史最优值
        self.gbestpop = np.zeros((self.dimpop, self.sizepop))  # 全局粒子历史最优位置(第二维复制）
        self.gbestval = np.zeros((1, 1)) + np.inf  # 全局粒子历史最优值

        self.detuning_frequency = detuning_frequency  # KPO的失谐频率
        self.kerr_constant = kerr_constant  # KPO的kerr常量
        self.pressure = pressure  # KPO泵涌压力
        self.time_step = time_step  # 时间步长
        self.symplectic_parameter = symplectic_parameter  # 辛欧拉方法参数
        self.convergence_threshold = convergence_threshold
        self.sampling_period = sampling_period

        self.xi0 = 0.7 * detuning_frequency / (np.std(ising.J) * ising.dimension ** (1 / 2))
        self.step = 0
        self.factor = 0
        self.symplectic_time_step = self.time_step / self.symplectic_parameter
        self.v_t = self.symplectic_time_step * self.detuning_frequency
        self.x_sample = np.sign(self.x)  # 粒子位置
        self.threshol_ratio = 0.9
        self.threshold = self.maxgen * self.threshol_ratio

        print(self.xi0)
    def v_update(self) -> None:

        """
        速度更新：
            lr[0]->sb算法动量更新\n
            lr[1]->个体最优更新\n
            lr[2]->群体最优更新

        Parameters
        ----
            self:
                sbpso pop

        Returns
        ----
            None
        """
        self.v -= self.symplectic_time_step * (self.kerr_constant * self.x ** 3 - self.factor * self.x)
        self.v += self.time_step * self.xi0 * (self.ising.J @ self.x)

    def x_update(self) -> None:

        """
        位置更新：
            新位置 = 原位置 + 速度 * 时间

        Parameters
        ----
            self:
                sbpso pop

        Returns
        ----
            None
        """

        self.x += self.v_t * self.v

    def factor_update(self) -> None:

        """
        时间参数

        Parameters
        ----
            self:
                sbpso pop

        Returns
        ----
            None
        """

        self.factor = self.pressure(self.step * self.time_step)

    def bestpop_update(self) -> None:

        """
        最佳粒子状态更新：
            个体最优更新\n
            群体最优更新

        Parameters
        ----
            self:
                sbpso pop

        Returns
        ----
            None
        """
        for i in range(self.sizepop):
            self.ising.get_ground_state(self.x[:, i, None])
            if self.pbestval[:, i] > self.ising.energy():
                self.pbestpop[:, i, None] = self.ising.ground_state.copy()
                self.pbestval[:, i] = self.ising.e.copy()
                print("当前粒子号：", i, "energy:", self.ising.e)
        if self.pbestval.min() < self.gbestval:
            for i in range(self.sizepop):
                if self.pbestval[:, i] == self.pbestval.min():
                    self.gbestpop = self.pbestpop[:, i, None].repeat(self.sizepop, axis=1)
                    self.gbestval = self.pbestval.min()

    def randomwalk(self) -> None:
        self.x = 0 * np.random.random((self.dimpop, self.sizepop))  # 粒子位置
        self.v = 0.01 * np.random.random((self.dimpop, self.sizepop))  # 粒子速度
        for self.step in range(self.maxgen):
            print(self.step, ":")
            self.factor_update()
            self.x_update()
            self.v_update()
            self.bestpop_update()
            if (self.step % 1000 == 0) or 1:
                print("gbestval is", self.gbestval)


gset_path = "./data/Gset/G22.txt"
maxcut = MaxCut.from_gset(gset_path)

ising_equivalent = maxcut.__to_Ising__()
a = sbpso(
    ising=ising_equivalent,
    weight=1,
    lr=(1, 0.1, 0.1),
    maxgen=3000,
    sizepop=20,
    rangepop=(-1, 1),
    rangespeed=(-0.5, 0.5)
)
print("The SB is processing...")
start_time = time.time()
a.randomwalk()
end_time = time.time()
for i in range(a.sizepop):
    a.ising.get_ground_state(a.x[:, i, None])
    maxcut.__from_Ising__(ising_equivalent)
print("final cut:", maxcut.cut_value)
print("time:", end_time - start_time, "s")
